import cv2
import numpy as np
import os
from skimage import graph, segmentation, color
from skimage.measure import regionprops
import shutil


def color_segmentation(
    image_path,
    k=8,
    min_pixels=0,
    color_thresh=15
):
    # 读取图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, _ = image.shape

    # 删除并重新创建输出目录
    if os.path.exists('./split'):
        shutil.rmtree('./split')  # 递归删除整个文件夹及其内容
    os.makedirs('./split', exist_ok=True)  # 重新创建空文件夹

    # K-means聚类
    pixel_data = image.reshape((-1, 3)).astype(np.float32)
    criteria = (cv2.TERM_CRITERIA_EPS +
                cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
    _, labels, centers = cv2.kmeans(
        pixel_data, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    centers = np.uint8(centers)
    clustered = centers[labels.flatten()].reshape((h, w, 3))
    labels = labels.reshape((h, w)) + 1  # 标签从1开始

    # 构建区域邻接图（RAG）并融合小色块
    if min_pixels > 0:
        rag = graph.rag_mean_color(image, labels)
        regions = regionprops(labels)

        # 标记需要融合的小色块
        small_regions = [r.label for r in regions if r.area < min_pixels]
        neighbor_map = {r.label: [] for r in regions}

        # 查找每个小色块的相邻色块
        for region in regions:
            if region.label in small_regions:
                boundary_labels = np.unique(labels[max(0, region.bbox[0]-1):min(h, region.bbox[2]+1),
                                                   max(0, region.bbox[1]-1):min(w, region.bbox[3]+1)])
                neighbor_map[region.label] = [
                    x for x in boundary_labels if x != region.label and x not in small_regions]

        # 融合到最相似的相邻色块
        for small_label in small_regions:
            if not neighbor_map[small_label]:
                continue
            # 计算颜色相似度
            small_color = rag.nodes[small_label]['mean color']
            min_diff = float('inf')
            best_neighbor = None
            for neighbor in neighbor_map[small_label]:
                neighbor_color = rag.nodes[neighbor]['mean color']
                diff = np.linalg.norm(small_color - neighbor_color)
                if diff < min_diff and diff < color_thresh:
                    min_diff = diff
                    best_neighbor = neighbor
            # 执行融合
            if best_neighbor:
                labels[labels == small_label] = best_neighbor

    # 保存色块（保持原分辨率）
    output_segments = []
    for label in np.unique(labels):
        if label == 0:
            continue  # 跳过背景
        mask = np.where(labels == label, 255, 0).astype(np.uint8)
        mean_color = cv2.mean(image, mask)[:3]

        # 创建带透明通道的色块
        segmented_img = np.zeros((h, w, 4), dtype=np.uint8)
        segmented_img[..., :3] = mean_color
        segmented_img[..., 3] = mask

        # 保存文件
        output_path = f'./split/segment_{label:03d}.png'
        cv2.imwrite(output_path, cv2.cvtColor(
            segmented_img, cv2.COLOR_RGBA2BGRA))
        output_segments.append(segmented_img)

    return output_segments, (h, w)

# 使用示例
if __name__ == "__main__":
    # 示例1：K-means聚类+小色块融合
    color_segmentation(
        './test.jpg',
        k=16,
        min_pixels=100,
        color_thresh=20
    )
